import scipy.io
import plotly.plotly as py
import plotly.graph_objs as go
import numpy as np
from scikits import audiolab
from matplotlib.pyplot import specgram
import matplotlib.pylab as plt
import glob
from multiprocessing import Pool
from scipy.signal.spectral import spectrogram
import pywt
from scipy.signal import find_peaks_cwt


def check_mat():
    py.sign_in('wuxiaominupc', '6jv5jrgvxs')
    mat = scipy.io.loadmat('/home/xiaomin/wxm/Data/KaggleSeizure2016/RawData/train_1/1_1_0.mat')
    print 1
    print mat.keys()
    print mat['dataStruct'][0, 0][0].shape
    print mat['dataStruct'][0, 0][0] - mat['dataStruct'][0][0][0]

    trace = go.Scatter(
        x=mat['dataStruct'][0][0][0][:10000, 0].transpose(),
        y=np.asarray([i for i in range(10000)])
    )
    data = [trace]
    py.plot(data, filename='KaggleSeizure2016Mat')


def test_plotly():
    N = 500
    py.sign_in('wuxiaominupc', '6jv5jrgvxs')
    Random_x = np.linspace(0, 1, N)
    Random_y = np.random.randn(N)
    trace = go.Scatter(
        x=Random_x,
        y=Random_y
    )
    data = [trace]
    py.plot(data, filename='KaggleSeizure2016Mat')


def test_kv():
    kv = {}
    for i in range(10):
        kv[str(i)] = str(i+1)
    print kv


def test_np_flat():
    a = np.random.randint(0, 1, (3, ))
    b = np.random.randint(0, 1, (2, 3))
    print (a.flat == b.flat).sum()


def check_pos_neg_ratio():
    lst = '/home/xiaomin/wxm/Data/KaggleSeizure2016/mats/train.lst'
    f = open(lst, 'r')
    labels = []
    num = len(open(lst, 'r').readlines())
    for i in range(num):
        index, path, label = f.readline().strip('\n').split('\t')
        labels.append(int(label))
    labels = np.array(labels)
    print labels
    print len(labels)
    print sum(labels)


def mat2wav():
    mat = scipy.io.loadmat('/home/xiaomin/wxm/Data/KaggleSeizure2016/RawData/train_1/1_1_0.mat')
    dat = mat['dataStruct'][0][0][0]

    mn = dat.min()
    mx = dat.max()
    mx = float(max(abs(mx), abs(mn)))
    if mx != 0:
        dat *= 0x7FFF / mx
    dat = np.int16(dat)
    for elec in range(16):
        print elec
        dstfile = '/home/xiaomin/wxm/Data/KaggleSeizure2016/tests/1_6_0_e' + str(elec) + '.wav'
        aud = dat[:, elec]
        audiolab.wavwrite(aud, dstfile, fs=400, enc='pcm16')


def test_wav_plot():
    wav_path = '/home/xiaomin/wxm/Data/KaggleSeizure2016/wavs/SubjPerChannel/train_1/1_6_0_e0.wav'
    wav = audiolab.wavread(wav_path)
    print wav[0]
    # Pxx, freqs, bins, im = specgram(wav[0], NFFT=512, Fs=400, noverlap=154, scale='dB')
    freqs, t, Pxx = spectrogram(wav[0], nperseg=512, fs=400, noverlap=154, window='hanning', mode='psd')
    print Pxx
    print Pxx.max()
    print Pxx.min()
    print Pxx.shape
    plt.show()


def test_glob_multiprocessing(index=0):
    folder = '/home/xiaomin/wxm/Data/KaggleSeizure2016/RawData/train_2'
    mat_paths = glob.glob(folder + '/' + '*.mat')
    print mat_paths


def check_wavelet():
    wav_path = '/home/xiaomin/wxm/Data/KaggleSeizure2016/wavs/SubjPerChannel/train_1/1_6_0_e0.wav'
    wav = audiolab.wavread(wav_path)
    print wav[0].shape
    mode = pywt.Modes.smooth
    # Pxx, freqs, bins, im = specgram(wav[0], NFFT=512, Fs=400, noverlap=154, scale='dB')
    # freqs, t, Pxx = spectrogram(wav[0], nperseg=512, fs=400, noverlap=154, window='hanning', mode='psd')
    w = pywt.Wavelet('coif5')
    (a, d) = pywt.dwt(wav[0], w, mode)
    print a.shape
    print d.shape


def find_peaks():
    wav_path = '/home/xiaomin/wxm/Data/KaggleSeizure2016/wavs/SubjPerChannel/train_1/1_6_0_e0.wav'
    wav = audiolab.wavread(wav_path)
    data = wav[0]
    x = np.linspace(0, len(data), num=len(data))
    fig = plt.figure()
    fig.subplots_adjust(hspace=0.2, bottom=.03, left=.07, right=.97, top=.92)
    ax = fig.add_subplot(2, 1, 1)
    ax.set_title("linchirp signal")
    ax.plot(x, data, 'b')
    ax.set_xlim(0, x[-1])

    ax = fig.add_subplot(2, 1, 2)
    ax.set_title("Wavelet packet coefficients at level %d" % 5)
    peaks = find_peaks_cwt(data, np.array([128, 512]))
    print peaks
    # ax.imshow()
    # ax.set_yticks(np.arange(0.5, len(labels) + 0.5), labels)

    plt.show()

if __name__ == '__main__':
    # p = Pool(4)
    # p.map(test_glob_multiprocessing, range(4))
    find_peaks()
